package com.ciaojian.core.filter;

import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.lang.Nullable;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.HandlerMapping;
import org.springframework.web.util.UrlPathHelper;

import javax.servlet.*;
import javax.servlet.annotation.WebFilter;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Locale;
import java.util.Set;

/**
 * spring jsession安全过滤器,防止 RDF 攻击
 *
 * @author Atlakyn
 */
@WebFilter(urlPatterns = "/*", filterName = "webFilter")
public class SpringJsessionidRdfFilter implements Filter {
    private final Set<String> safeExtensions = new HashSet<>();
    /**
     * Extensions associated with the built-in message converters
     */
    private static final Set<String> WHITELISTED_EXTENSIONS = new HashSet<>(Arrays.asList(
            "txt", "text", "yml", "properties", "csv",
            "json", "xml", "atom", "rss",
            "png", "jpe", "jpeg", "jpg", "gif", "wbmp", "bmp"));

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        HttpServletRequest request = (HttpServletRequest) servletRequest;
        HttpServletResponse response = (HttpServletResponse) servletResponse;

        String contentDisposition = response.getHeader(HttpHeaders.CONTENT_DISPOSITION);
        if (!"".equals(contentDisposition) && null != contentDisposition) {
            return;
        }

        try {
            int status = response.getStatus();
            if (status < 200 || status > 299) {
                return;
            }
        } catch (Throwable ex) {
            // ignore
        }

        String requestUri = request.getRequestURI();

        System.out.println(requestUri);

        if (requestUri.contains(";jsessionid=")) {
            int index = requestUri.lastIndexOf('/') + 1;
            String filename = requestUri.substring(index);
            String pathParams = "";

            index = filename.indexOf(';');
            if (index != -1) {
                pathParams = filename.substring(index);
                filename = filename.substring(0, index);
            }

            UrlPathHelper decodingUrlPathHelper = new UrlPathHelper();
            filename = decodingUrlPathHelper.decodeRequestString(request, filename);
            String ext = StringUtils.getFilenameExtension(filename);

            pathParams = decodingUrlPathHelper.decodeRequestString(request, pathParams);
            String extInPathParams = StringUtils.getFilenameExtension(pathParams);

            if (!safeExtension(request, ext) || !safeExtension(request, extInPathParams)) {
                response.addHeader(HttpHeaders.CONTENT_DISPOSITION, "inline;filename=f.txt");
            }
        }
        filterChain.doFilter(servletRequest, servletResponse);
    }

    private boolean safeExtension(HttpServletRequest request, @Nullable String extension) {
        if (!StringUtils.hasText(extension)) {
            return true;
        }
        extension = extension.toLowerCase(Locale.ENGLISH);
        this.safeExtensions.addAll(WHITELISTED_EXTENSIONS);
        if (this.safeExtensions.contains(extension)) {
            return true;
        }
        String pattern = (String) request.getAttribute(HandlerMapping.BEST_MATCHING_PATTERN_ATTRIBUTE);
        if (pattern != null && pattern.endsWith("." + extension)) {
            return true;
        }
        if (extension.equals("html")) {
            String name = HandlerMapping.PRODUCIBLE_MEDIA_TYPES_ATTRIBUTE;
            Set<MediaType> mediaTypes = (Set<MediaType>) request.getAttribute(name);
            if (!CollectionUtils.isEmpty(mediaTypes) && mediaTypes.contains(MediaType.TEXT_HTML)) {
                return true;
            }
        }
        return false;
    }

}
